import os
import glob
from PIL import Image

def merge_images(input_dir, output_path, n_rows=5, n_cols=2):
    # 获取所有png文件
    png_files = glob.glob(os.path.join(input_dir, '*_assim_plot.png'))  # 只获取同化结果图
    
    if not png_files:
        print("未找到PNG文件")
        return
        
    # 按照像素ID排序
    png_files.sort(key=lambda x: int(os.path.basename(x).split('_')[0]))
    
    # 打开第一张图片获取尺寸
    first_img = Image.open(png_files[0])
    img_width, img_height = first_img.size
    
    # 创建新图片
    merged_width = n_cols * img_width
    merged_height = n_rows * img_height
    merged_img = Image.new('RGB', (merged_width, merged_height), 'white')
    
    # 拼接图片
    for idx, png_file in enumerate(png_files[:n_rows * n_cols]):  # 限制最多10张图片
        img = Image.open(png_file)
        row = idx // n_cols  # 计算行号
        col = idx % n_cols   # 计算列号
        merged_img.paste(img, (col * img_width, row * img_height))
        
    # 保存结果
    merged_img.save(output_path)
    print(f"图片已保存至: {output_path}")

def merge_images_3x3(input_dir, output_path, n_rows=3, n_cols=3):
    # 获取所有png文件
    png_files = glob.glob(os.path.join(input_dir, '*_assim_plot.png'))
    
    if not png_files:
        print("未找到PNG文件")
        return
        
    # 按照像素ID排序
    png_files.sort(key=lambda x: int(os.path.basename(x).split('_')[0]))
    
    # 打开第一张图片获取尺寸
    first_img = Image.open(png_files[0])
    img_width, img_height = first_img.size
    
    # 创建新图片
    merged_width = n_cols * img_width
    merged_height = n_rows * img_height
    merged_img = Image.new('RGB', (merged_width, merged_height), 'white')
    
    # 拼接前9张图片
    for idx, png_file in enumerate(png_files[:n_rows * n_cols]):
        img = Image.open(png_file)
        row = idx // n_cols
        col = idx % n_cols
        merged_img.paste(img, (col * img_width, row * img_height))
        
    # 保存结果
    output_path_3x3 = os.path.join(os.path.dirname(output_path), "merged_3x3.png")
    merged_img.save(output_path_3x3)
    print(f"3x3拼接图已保存至: {output_path_3x3}")

if __name__ == "__main__":
    input_dir = r"D:\code\MAS_DA\output_agents_fixed_v2\png"
    output_path = os.path.join(input_dir, "merged3x3.png")
    
    # 生成5*2的拼接图
    # merge_images(input_dir, output_path, n_rows=5, n_cols=2)
    
    # 生成3*3的拼接图
    merge_images_3x3(input_dir, output_path)
